{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/pytorch/use_nano_decorator_pytorch_training.ipynb)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use `@nano` Decorator to Accelerate PyTorch Training Loop"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "BigDL-Nano integrates multiple optimizations to accelerate PyTorch training workloads. As a pure PyTorch user, you could simply wrap your custom PyTorch training loop with `@nano` decorator to benefit from BigDL-Nano."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "To use `@nano` decorator, you need to install BigDL-Nano for PyTorch first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "!pip install --pre --upgrade bigdl-nano[pytorch] # install the nightly-built version\n",
    "!source bigdl-nano-init # set environment variables"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> Before starting your PyTorch application, it is highly recommended to run `source bigdl-nano-init` to set several environment variables based on your current hardware. Empirically, these variables will bring big performance increase for most PyTorch applications on training workloads."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "> ⚠️ **Warning**\n",
    "> \n",
    "> For Jupyter Notebook users, we recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose you define your custom PyTorch training loop as follows. To benefit from BigDL-Nano integrated optimizations, you could simply **import** `nano` **decorator, and wrap the training loop with it.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from bigdl.nano.pytorch import nano # import nano decorator\n",
    "\n",
    "@nano() # apply the decorator to the training loop\n",
    "def training_loop(model, optimizer, train_loader, num_epochs, loss_func):\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "\n",
    "        model.train()\n",
    "        train_loss, num = 0, 0\n",
    "        with tqdm(train_loader, unit=\"batch\") as tepoch:\n",
    "            for data, target in tepoch:\n",
    "                tepoch.set_description(f\"Epoch {epoch}\")\n",
    "                optimizer.zero_grad()\n",
    "                output = model(data)\n",
    "                loss = loss_func(output, target)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                loss_value = loss.sum()               \n",
    "                train_loss += loss_value\n",
    "                num += 1\n",
    "                tepoch.set_postfix(loss=loss_value)\n",
    "            print(f'Train Epoch: {epoch}, avg_loss: {train_loss / num}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> To make sure `@nano` is functional on your custom training loop, there are some requirements for its **parameter lists**:\n",
    ">\n",
    "> - there should be one and only one instance of `torch.nn.Module` passed in the training loop as model \n",
    "> - there should be at least one instance of `torch.optim.Optimizer` passed in the training loop as optimizer\n",
    "> - there should be at least one instance of `torch.utils.data.DataLoader` passed in the training loop as dataloader"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could then call the `training_loop` method as normal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# Define model and dataloader\n",
    "\n",
    "from torch import nn\n",
    "from torchvision.models import resnet18\n",
    "\n",
    "class MyPytorchModule(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = resnet18(pretrained=True)\n",
    "        num_ftrs = self.model.fc.in_features\n",
    "        # Here the size of each output sample is set to 37.\n",
    "        self.model.fc = nn.Linear(num_ftrs, 37)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import OxfordIIITPet\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "def create_train_dataloader():\n",
    "    train_transform = transforms.Compose([transforms.Resize(256),\n",
    "                                          transforms.RandomCrop(224),\n",
    "                                          transforms.RandomHorizontalFlip(),\n",
    "                                          transforms.ColorJitter(brightness=.5, hue=.3),\n",
    "                                          transforms.ToTensor(),\n",
    "                                          transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                               [0.229, 0.224, 0.225])])\n",
    "\n",
    "    # Apply data augmentation to the tarin_dataset\n",
    "    train_dataset = OxfordIIITPet(root=\"/tmp/data\", transform=train_transform, download=True)\n",
    "\n",
    "    # prepare train data loaders\n",
    "    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "    return train_dataloader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MyPytorchModule()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "loss_func = torch.nn.CrossEntropyLoss()\n",
    "train_loader = create_train_dataloader()\n",
    "\n",
    "training_loop(model, optimizer, train_loader, num_epochs=5, loss_func=loss_func)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; _The definition of_ `MyPytorchModule` _and_ `create_train_dataloader` _can be found in the_ [runnable example](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/pytorch/use_nano_decorator_pytorch_training.ipynb)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> Due to the optimized environment variables set by `source bigdl-nano-init`, you could already experience some training acceleration after wrapping your custom training loop with `@nano` decorator.\n",
    "> \n",
    "> For more optimizations provided by `@nano` decorator, you can refer to the Related Readings."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Readings**\n",
    "> \n",
    "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)\n",
    "> - [How to accelerate a PyTorch application on training workloads through Intel® Extension for PyTorch*](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_ipex.html)\n",
    "> - [How to accelerate a PyTorch application on training workloads through multiple instances](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_multi_instance.html)\n",
    "> - [How to use the channels last memory format in your PyTorch application for training](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/pytorch_training_channels_last.html)\n",
    "> - [How to conduct BFloat16 Mixed Precision training in your PyTorch application](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_bf16.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nano-pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.15 (default, Nov 24 2022, 21:12:53) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "09344c7f3239fd422839751f876786d6b1a624c40f19af1b43cb2737f421c2b2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
